import tvm
import tvm.relay as relay
from tvm.relay import testing
from tvm.relax.testing import relay_translator

MODEL_SEEDS = [
    # get_lstm(),
    testing.resnet.get_workload(batch_size=1, num_layers=18, image_shape=(128, 128, 3), layout="NHWC"),
    testing.squeezenet.get_workload(batch_size=1, num_classes=100, image_shape=(3, 128, 128), dtype='float32'),
    testing.mobilenet.get_workload(image_shape=(3, 128, 128)),
    testing.mlp.get_workload(batch_size=1, num_classes=10, image_shape=(1, 64, 64)),
    testing.dcgan.get_workload(batch_size=1),
    testing.inception_v3.get_workload(),
    testing.vgg.get_workload(batch_size=1),
    testing.densenet.get_workload(),
]

if __name__ == '__main__':
    for i, (mod, params) in enumerate(MODEL_SEEDS):
         relax_mod = relay_translator.from_relay(mod['main'], tvm.target.Target("llvm"))       
         file_name = f'tzer_seeds/{i}.txt'
         with open(file_name, 'w') as f:
             str_relax_mod = str(relax_mod)
             f.write(str_relax_mod)
         with open(file_name, 'r') as f:
             relax_mod_str = f.read()
             new_ir = tvm.script.from_source(relax_mod_str)
             print(new_ir)
         print('yes')
